#!/usr/bin/env python3
# coding=utf-8
# date 2024-05-20
# author calllivecn <calllivecn@outlook.com>
# Windows 客户端适配

import io
import os
import sys
import enum
import struct
import socket
import asyncio
import logging
import argparse
import ctypes
from ctypes import wintypes
import threading
from queue import Queue

from logging.handlers import TimedRotatingFileHandler

from asyncio import (
    StreamReader,
    StreamWriter,
)

from typing import (
    Union,
    Optional,
    BinaryIO,
    Tuple,
    TypeVar,
)

Self = TypeVar("Self")

BUFSIZE = 1<<12 # 4K

def getlogger(level=logging.INFO):
    fmt = logging.Formatter("%(asctime)s.%(msecs)03d %(levelname)s %(filename)s:%(funcName)s:%(lineno)d %(message)s", datefmt="%Y-%m-%d-%H:%M:%S")

    prog, _ = os.path.splitext(os.path.basename(sys.argv[0]))
    fp = TimedRotatingFileHandler(f"{prog}.logs", when="D", interval=1, backupCount=7)
    fp.setFormatter(fmt)

    logger = logging.getLogger(f"{prog}")
    logger.setLevel(level)
    logger.addHandler(fp)
    return logger

logger = getlogger()

# Windows API 常量和函数
KERNEL32 = ctypes.windll.kernel32
STD_INPUT_HANDLE = -10
STD_OUTPUT_HANDLE = -11

# CONSOLE_SCREEN_BUFFER_INFO 结构
class CONSOLE_SCREEN_BUFFER_INFO(ctypes.Structure):
    _fields_ = [
        ("dwSize", wintypes._COORD),
        ("dwCursorPosition", wintypes._COORD),
        ("wAttributes", wintypes.WORD),
        ("srWindow", wintypes.SMALL_RECT),
        ("dwMaximumWindowSize", wintypes._COORD),
    ]

def get_console_size() -> Tuple[int, int]:
    """获取控制台窗口大小 (列, 行)"""
    h_stdout = KERNEL32.GetStdHandle(STD_OUTPUT_HANDLE)
    csbi = CONSOLE_SCREEN_BUFFER_INFO()
    KERNEL32.GetConsoleScreenBufferInfo(h_stdout, ctypes.byref(csbi))
    columns = csbi.srWindow.Right - csbi.srWindow.Left + 1
    rows = csbi.srWindow.Bottom - csbi.srWindow.Top + 1
    return rows, columns
    # return columns, rows

def set_console_size(columns: int, rows: int):
    """设置控制台窗口大小"""
    h_stdout = KERNEL32.GetStdHandle(STD_OUTPUT_HANDLE)
    # 设置控制台缓冲区大小
    buf_size = wintypes._COORD(columns, rows)
    KERNEL32.SetConsoleScreenBufferSize(h_stdout, buf_size)
    # 设置控制台窗口大小
    rect = wintypes.SMALL_RECT(0, 0, columns-1, rows-1)
    KERNEL32.SetConsoleWindowInfo(h_stdout, True, ctypes.byref(rect))

# 定义Windows上的终端原始模式设置
class TerminalMode:
    def __init__(self):
        self.h_stdin = KERNEL32.GetStdHandle(STD_INPUT_HANDLE)
        self.old_mode = wintypes.DWORD()
        self.success = KERNEL32.GetConsoleMode(self.h_stdin, ctypes.byref(self.old_mode))
    
    def set_raw_mode(self):
        """设置终端为原始模式"""
        if not self.success:
            return
            
        # 原始模式下的控制台模式
        # 禁用以下模式：
        # 0x0001 - 启用行输入
        # 0x0002 - 启用回显
        # 0x0004 - 启用插入模式
        # 0x0008 - 启用自动换行
        # 0x0010 - 启用快速编辑模式
        # 0x0020 - 启用扩展属性
        # 0x0040 - 启用Processed Output
        # 0x0080 - 启用Virtual Terminal Input
        # 0x0100 - 启用LCTRL+C Break Handling
        raw_mode = self.old_mode.value & ~(0x0001 | 0x0002 | 0x0004 | 0x0008 | 0x0010 | 0x0020 | 0x0040 | 0x0080 | 0x0100)
        KERNEL32.SetConsoleMode(self.h_stdin, raw_mode)
    
    def restore(self):
        """恢复原始终端模式"""
        if self.success:
            KERNEL32.SetConsoleMode(self.h_stdin, self.old_mode.value)

TermSize = struct.Struct("HH")

class PacketError(Exception):
    pass

class PacketType(enum.IntEnum):
    """
    一字节，数据包类型。
    """
    ZERO = 0 # 保留
    EXIT = enum.auto()
    TTY_RESIZE = enum.auto()
    TRANER = enum.auto()

class Packet:
    """
    1B: protocol version
    1B: package type
    2B: payload lenght
    data: payload
    """

    header = struct.Struct("!BBH")

    hsize = header.size

    def __init__(self):
        self.Version = 0x01 # protocol version

    def tobuf(self, typ: PacketType, data: bytes) -> Union[bytes, memoryview]:
        self.buf = io.BytesIO()
        h = self.header.pack(self.Version, typ, len(data))
        self.buf.write(h)
        self.buf.write(data)
        return self.buf.getvalue()
    
    def frombuf(self, data: bytes):
        self.Version, typ, lenght = self.header.unpack(data[:self.hsize])
        self.typ = PacketType(typ)
        self.lenght = lenght

class RecvSend:

    def __init__(self, reader: StreamReader, writer: StreamWriter):
        self.reader = reader
        self.writer = writer

        self.max_payload = 65536
    
    async def read(self) -> Tuple[PacketType, Union[bytes, memoryview]]:
        logger.debug(f"__read() Packet.hsize")

        data = await self.__read(Packet.hsize)
        if data == b"":
            return PacketType.EXIT, b""

        if len(data) != Packet.hsize:
            raise PacketError("invalid data")
        
        ph = Packet()
        ph.frombuf(data)

        payload = await self.__read(ph.lenght)

        logger.debug(f"__read() typ:{ph.typ.name}, lenght:{ph.lenght} payload:{payload}")


        if len(payload) != ph.lenght:
            raise PacketError("invalid data")
        
        return ph.typ, payload
    
    async def write(self, typ: PacketType, data: bytes) -> int:
        data_len = len(data)
        if data_len > self.max_payload:
            raise ValueError(f"数据包太大： 0 ~ {self.max_payload}")
        
        ph = Packet()
        payload = ph.tobuf(typ, data)

        await self.__write(payload)

    
    def getsockname(self):
        return self.writer.get_extra_info("peername")

    async def close(self):
        self.writer.close()
        await self.writer.wait_closed()
    
    async def __read(self, size: int) -> Union[bytes, memoryview]:
        buf = io.BytesIO()
        while (d := await self.reader.read(size)) != b"":
            buf.write(d)
            size -= len(d)
        
        return buf.getvalue()
    
    async def __write(self, data: Union[bytes, memoryview]):
        self.writer.write(data)
        await self.writer.drain()
        return len(data)

# 同步读取stdin并放入队列
class StdinReader:
    def __init__(self, loop):
        self.loop = loop
        self.queue = Queue(maxsize=100)
        self.stopped = False
    
    def read_stdin(self):
        try:
            while not self.stopped:
                # 同步读取stdin
                data = os.read(sys.stdin.fileno(), BUFSIZE)
                if not data:
                    break
                # 将数据放入队列并通知asyncio
                self.loop.call_soon_threadsafe(self.queue.put_nowait, data)
        except Exception as e:
            logger.error(f"read_stdin error: {e}")
        finally:
            self.loop.call_soon_threadsafe(self.queue.put_nowait, b"")

# client 从队列 -> sock
async def stdin2sock(w: RecvSend, stdin_reader: StdinReader):
    try:
        while True:
            payload = await asyncio.to_thread(stdin_reader.queue.get)
            if not payload:
                break
            await w.write(PacketType.TRANER, payload)
            logger.debug(f"stdin2sock: {payload}")
    except asyncio.CancelledError:
        logger.debug("stdin2sock Cancelled")

    logger.debug("stdin2sock exit.")

async def sock2stdout(r: RecvSend):
    while True:
        logger.debug(f"read()")
        typ, payload = await r.read()
        logger.debug(f"{PacketType(typ).name}, {payload}")

        if typ == PacketType.TRANER:
            logger.debug(f"sys.stdout.write(payload) --> stdout")
            try:
                # 尝试以UTF-8解码并写入
                sys.stdout.write(payload.decode('utf-8'))
            except UnicodeDecodeError:
                # 如果解码失败，将其作为二进制数据写入
                sys.stdout.buffer.write(payload)
            sys.stdout.flush()

        elif typ == PacketType.TTY_RESIZE:
            columns, rows = TermSize.unpack(payload)
            logger.debug(f"tty resize: columns={columns}, rows={rows}")
            set_console_size(columns, rows)

        elif typ == PacketType.EXIT:
            await r.close()
            logger.debug(f"peer sock close")
            break
        else:
            logger.debug(f"不符合的包：{typ}, {payload}")

    logger.debug(f"done")

async def handle_window_resize(traner: RecvSend):
    """处理窗口大小变化"""
    h_stdin = KERNEL32.GetStdHandle(STD_INPUT_HANDLE)
    
    # 定义 INPUT_RECORD 结构
    class INPUT_RECORD(ctypes.Structure):
        _fields_ = [
            ("EventType", wintypes.WORD),
            ("Event", ctypes.c_ubyte * 28),  # 占位符，实际使用时需根据 EventType 解析
        ]
    
    INPUT_RECORD_SIZE = ctypes.sizeof(INPUT_RECORD)
    BUFFER_SIZE = 10
    input_buffer = (INPUT_RECORD * BUFFER_SIZE)()
    events_read = wintypes.DWORD()
    
    last_size = get_console_size()
    
    while True:
        try:
            # 异步读取控制台输入事件
            await asyncio.sleep(1)  # 避免占用过多 CPU
            
            # 非阻塞读取控制台输入事件
            if KERNEL32.PeekConsoleInputA(h_stdin, ctypes.byref(input_buffer), BUFFER_SIZE, ctypes.byref(events_read)):
                for i in range(events_read.value):
                    event = input_buffer[i]
                    
                    # 检查窗口缓冲区大小变化事件
                    if event.EventType == 0x0010:  # WINDOW_BUFFER_SIZE_EVENT
                        new_size = get_console_size()
                        if new_size != last_size:
                            print(f"recv window resize: {new_size}")
                            await traner.write(PacketType.TTY_RESIZE, TermSize.pack(*new_size))
                            last_size = new_size
                    
                    # 其他事件可以忽略，但需要从输入缓冲区移除
                    KERNEL32.ReadConsoleInputA(h_stdin, ctypes.byref(input_buffer), BUFFER_SIZE, ctypes.byref(events_read))
                    
        except Exception as e:
            logger.error(f"处理窗口大小变化出错: {e}")
            break

async def client(args):
    """Windows 客户端实现"""
    r, w = await asyncio.open_connection(args.addr, args.port)
    print('connected to server')
    traner = RecvSend(r, w)

    loop = asyncio.get_running_loop()

    # 初始化对面终端
    columns, rows = get_console_size()
    print(f"send my terminal size: {columns}x{rows}")
    await traner.write(PacketType.TTY_RESIZE, TermSize.pack(columns, rows))

    # 设置终端为原始模式
    term_mode = TerminalMode()
    term_mode.set_raw_mode()

    try:
        # 创建同步stdin阅读器
        stdin_reader = StdinReader(loop)
        
        # 启动stdin读取线程
        stdin_thread = threading.Thread(target=stdin_reader.read_stdin)
        stdin_thread.daemon = True
        stdin_thread.start()

        # 启动任务
        stdin2sock_task = asyncio.create_task(stdin2sock(traner, stdin_reader))
        sock2stdout_task = asyncio.create_task(sock2stdout(traner))
        resize_task = asyncio.create_task(handle_window_resize(traner))

        await sock2stdout_task
        
    finally:
        # 停止stdin读取线程
        stdin_reader.stopped = True
        # 恢复终端模式
        term_mode.restore()
        logger.debug("restore termios")
        await traner.close()

def main():
    parse = argparse.ArgumentParser(
        usage="%(prog)s [option]",
        description="远程shell连接 Windows 客户端",
        epilog="END",
    )

    parse.add_argument("--debug", action="store_true", help="debug")
    parse.add_argument("--addr", default="::", help="需要连接的IP或域名")
    parse.add_argument("--port", default=6789, type=int, help="端口(default: 6789")

    args = parse.parse_args()

    if args.debug:
        logger.setLevel(logging.DEBUG)
    
    asyncio.run(client(args), debug=True)

if __name__ == "__main__":
    main()
